#!/bin/bash
# =============================================================================
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================

export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

# 此处用于设置动态链接库搜索路径，请根据当前机器的实际路径设置
CPP_BASE_DIR="/usr/include/c++"
CPP_VERSION=$(ls $CPP_BASE_DIR | grep -E '^[0-9]+(\.[0-9]+)*$' | sort -V | tail -n 1)

if [ -z "$CPP_VERSION" ]; then
    echo "No valid C++ standard library version found in $CPP_BASE_DIR."
else
    echo "Detected C++ standard library version: $CPP_VERSION"
    export CPLUS_INCLUDE_PATH="$CPP_BASE_DIR/$CPP_VERSION:$CPP_BASE_DIR/$CPP_VERSION/aarch64-linux-gnu:$CPP_BASE_DIR/$CPP_VERSION/aarch64-linux-gnu/bits:$CPLUS_INCLUDE_PATH"
    echo "CPLUS_INCLUDE_PATH set to: $CPLUS_INCLUDE_PATH"
fi


SCRIPT_DIR=$(dirname "$(realpath "$0")")
PROJECT_DIR=$(realpath "$SCRIPT_DIR"/../)
export PYTHONPATH=${PYTHONPATH}:${PROJECT_DIR}
export IGNORE_INFER_ERROR=1
echo SCRIPT_DIR=$SCRIPT_DIR
echo PROJECT_DIR=$PROJECT_DIR

CKPT_PATH="/PATH/TO/CHECKPOINT"
FAST_CKPT_PATH=="/PATH/TO/CACHE_CHECKPOINT"

CMD="torchrun --nproc_per_node=8 --nnodes=1 \
        "${PROJECT_DIR}/src/generate.py" \
        --input-texts \
        --global-batch-size 1 \
        --en-sequence-len 2048 \
        --max-new-tokens 128 \
        --streaming \
        --atten-tp-size 8 \
        --atten-dp-size 1 \
        --routed-expert-tp-size 1 \
        --routed-expert-ep-size 8 \
        --share-expert-tp-size 8 \
        --share-expert-ep-size 1 \
        --H2P \
        --compile \

        --ckpt-dtype float16 \
        --ckpt-path $CKPT_PATH \
        --fast-ckpt-path $FAST_CKPT_PATH \
        "


# 可选参数（取消注释以启用）
# CMD="${CMD} --temperature 0.6"
# CMD="${CMD} --top-k 50"
# CMD="${CMD} --top-p 0.95"
# CMD="${CMD} --fast-ckpt-path ${CKPT_PATH_FAST}"
# CMD="${CMD} --cache-compile-path ${CACHE_COMPILE_PATH}"
# CMD="${CMD} --input-file ${INPUT_FILE}"


echo ${CMD} | sed "s/--/\n  --/g"
echo "============================================"

${CMD} 2>&1 | tee ${PROJECT_DIR}/inference.log


